from ultralytics import YOLO
import cv2
import numpy as np
import argparse

# 创建 ArgumentParser 对象
parser = argparse.ArgumentParser(description='输入输出文件')

# 添加参数
parser.add_argument('-f', help='输入文件')
parser.add_argument('-d', help='输出文件', default='default value')

model = YOLO("yolo11m-seg.pt")

def inference(image):
    return model(image)

def draw_result_image(image, results):
    print(results)
    segmentation_mask = np.zeros_like(image, dtype=np.uint8)
    # Iterate over the results
    for i, r in enumerate(results):
        # Iterate through the detected masks
        for j, mask in enumerate(r.masks.xy):
            # Convert the class tensor to an integer
            class_id = int(r.boxes.cls[j].item())  # Extract the class ID as an integer
            
            # Check if the detected class corresponds to 'person' (class ID 0)
            # if class_id == 0:
            # Convert mask coordinates to an integer format for drawing
            mask = np.array(mask, dtype=np.int32)
            
            # Fill the segmentation mask with color
            cv2.fillPoly(segmentation_mask, [mask], (0, 255, 0))

    # Combine the original image with the segmentation mask
    segmentation_result = cv2.addWeighted(image, 1, segmentation_mask, 0.7, 0)

    # Optionally display the image (make sure you're running in a GUI environment)
    cv2.imshow("Segmentation Result", segmentation_result)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

def with_video():
    cap = cv2.VideoCapture(0)
    while True:
        success, img = cap.read()
        if not success:
            break
        results = inference(img)
        draw_result_image(img, results)

def with_image(filename):
    results = inference(filename)
    draw_result_image(img, results)

def gen_seg(filein, fileout):
    image = cv2.imread(filein)
    results = inference(filein)
    draw_result_image(image, results)

if __name__ == "__main__":
    args = parser.parse_args()
    gen_seg(args.f, "")

# python yolov11-seg.py -f /Users/v/Documents/proj/testdata/origin_imgs/images_labelimg_yolo/IMG_20181201_181804.jpg
